﻿#pragma once
#include<iostream>
#include<assert.h>
#include<vector>
using namespace std;
template<class K,class V>
struct TreeNode
{
	pair<K, V> _kv;
	TreeNode<K, V>* _left;
	TreeNode<K, V>* _right;
	TreeNode<K, V>* _parent;
	int _bf;
	TreeNode(const pair<K, V>& kv)
		:_kv(kv),
		_left(nullptr),
		_right(nullptr),
		_parent(nullptr),
		_bf(0) {
		;
	}
};
template<class K,class V>
class AVLTree {
public:
	using TNode = TreeNode<K, V>;

	//RR型 ：只要左旋
	void Rotation_Left(TNode* unbalancd_node, TNode* unbalanced_node_right) {
		if (unbalancd_node == _root) {
			unbalancd_node->_right = unbalanced_node_right->_left;
			if (unbalanced_node_right->_left!=nullptr) {
				unbalanced_node_right->_left->_parent = unbalancd_node;
			}
			unbalanced_node_right->_left = unbalancd_node;
			_root = unbalanced_node_right;
			_root->_parent = nullptr;
			unbalancd_node->_parent = _root;

			//更新平衡因子：
			unbalancd_node->_bf = 0;
			unbalanced_node_right->_bf = 0;
			return;
		}

		TNode* unbalanced_parent = unbalancd_node->_parent;
		if (unbalanced_parent->_left == unbalancd_node) { //不平衡节点在其父节点的左边
			unbalanced_parent->_left = unbalanced_node_right;
			unbalanced_node_right->_parent = unbalanced_parent; 
		}
		else {//不平衡节点在其父节点的右边
			unbalanced_parent->_right = unbalanced_node_right;
			unbalanced_node_right->_parent = unbalanced_parent;
		}


		unbalancd_node->_right = unbalanced_node_right->_left;
		if (unbalanced_node_right->_left != nullptr) {
			unbalanced_node_right->_left->_parent = unbalancd_node;
		}
		unbalanced_node_right->_left = unbalancd_node;
		unbalancd_node->_parent = unbalanced_node_right;

		//更新平衡因子：
		unbalancd_node->_bf = 0;
		unbalanced_node_right->_bf = 0;
	}

	//LL型：只要右旋
	void Rotation_Right(TNode* unbalancd_node, TNode* unbalanced_node_left) {
		if (unbalancd_node == _root) {
			unbalancd_node->_left = unbalanced_node_left->_right;
			if (unbalanced_node_left->_right != nullptr) {
				unbalanced_node_left->_right->_parent = unbalancd_node;
			}
			unbalanced_node_left->_right = unbalancd_node;
			_root = unbalanced_node_left;
			_root->_parent = nullptr;
			unbalancd_node->_parent = _root;

			//更新平衡因子：
			unbalancd_node->_bf = 0;
			unbalanced_node_left->_bf = 0;
			return;
		 }
		TNode* unbalanced_parent = unbalancd_node->_parent;
		if (unbalancd_node == unbalanced_parent->_left) { //不平衡节点在其父节点的左边
			unbalanced_parent->_left = unbalanced_node_left;
			unbalanced_node_left->_parent = unbalanced_parent;
		}
		else {//不平衡节点在其父节点的右边
			unbalanced_parent->_right = unbalanced_node_left;
			unbalanced_node_left->_parent = unbalanced_parent;
		}

		unbalancd_node->_left = unbalanced_node_left->_right;
		if (unbalanced_node_left->_right != nullptr) {
			unbalanced_node_left->_right->_parent = unbalancd_node;
		}
		unbalanced_node_left->_right = unbalancd_node;
		unbalancd_node->_parent = unbalanced_node_left;

		//更新平衡因子：
		unbalancd_node->_bf = 0;
		unbalanced_node_left->_bf = 0;
	}

	//LR型：先左单旋再右单旋
	void Rotation_LR(TNode* unbalancd_node, TNode* unbalanced_node_left) {
		TNode* unbalanced_node_left_right = unbalanced_node_left->_right;
		//提前保存unbalanced_node_left_right的平衡因子
		int _bf = unbalanced_node_left_right->_bf;

		Rotation_Left(unbalanced_node_left, unbalanced_node_left_right);
		Rotation_Right(unbalancd_node, unbalanced_node_left_right);

		//更新平衡因子：需要更新平衡因子的节点只有三个

		if (_bf == 0) { //说明unb_L_R节点是新插入的
			unbalancd_node->_bf = 0;
			unbalanced_node_left->_bf = 0;
			unbalanced_node_left_right->_bf = 0;
		}
		else if (_bf == 1) {//说明新节点插入在unb_L_R的右边
			unbalancd_node->_bf = 0;
			unbalanced_node_left->_bf = -1;
			unbalanced_node_left_right->_bf = 0;
		}
		else if (_bf == -1) { //说明新节点插入在unb_L_R的左边
			unbalancd_node->_bf = 1;
			unbalanced_node_left->_bf = 0;
			unbalanced_node_left_right->_bf = 0;
		}
		else {
			assert(false); 
		}
	}

	//RL型：先右单旋再左单旋
	void Rotation_RL(TNode* unbalancd_node, TNode* unbalanced_node_right) {
		TNode* unbalanced_node_right_left = unbalanced_node_right->_left;
		//提前记录unb_r_l的平衡因子
		int _bf = unbalanced_node_right_left->_bf;

		Rotation_Right(unbalanced_node_right, unbalanced_node_right_left);
		Rotation_Left(unbalancd_node, unbalanced_node_right_left);

		//更新平衡因子：
		if (_bf == 0) { //说明unb_r_l节点是新插入的
			unbalancd_node->_bf = 0;
			unbalanced_node_right->_bf = 0;
			unbalanced_node_right_left->_bf = 0;
		}
		else if (_bf == 1) { //说明新节点插入在unb_r_l的右边
			unbalancd_node->_bf = -1;
			unbalanced_node_right->_bf = 0;
			unbalanced_node_right_left->_bf = 0;
		}
		else if (_bf == -1) { //说明新节点插入在unb_r_l的左边
			unbalancd_node->_bf = 0;
			unbalanced_node_right->_bf = 1;
			unbalanced_node_right_left->_bf = 0;
		}
		else {
			assert(false);
		}
	}


	bool Insert(const pair<K, V>& kv) {
		if (_root == nullptr) {
			_root = new TNode(kv);
			return true;
		}

		//第一步：插入
		//逻辑与普通二叉搜索树一致

		TNode* cur = _root;
		TNode* parent = nullptr;
		while (cur) {
			if (cur->_kv.first > kv.first) {
				parent = cur;
				cur = cur->_left;
			}
			else if (cur->_kv.first < kv.first) {
				parent = cur;
				cur = cur->_right;
			}
			else {
				return false;
			}
		}
		TNode* newnode = new TNode(kv);
		if (parent->_kv.first > kv.first) {
			parent->_left = newnode;
			newnode->_parent = parent;
		}
		else {
			parent->_right = newnode;
			newnode->_parent = parent;
		}

		TNode* src = newnode;
		while (parent) {
			if (parent->_left == src)
				parent->_bf--;
			else
				parent->_bf++;
			if (parent->_bf == 0) {
				break;
			}
			else if (parent->_bf == 1 || parent->_bf == -1) { //继续更新
				src = parent;
				parent = parent->_parent;
			}

			else { //parent->_bf == 2 || parent->_bf == -2 ：需要旋转
				src = newnode;
				//第三步：旋转	
				if (parent->_bf == 0) {
					break;
				}
	         	else if (parent->_bf == 2) {
					TNode* unbalanced_node = parent; //不平衡节点
					TNode* unbalanced_node_right = parent->_right; //不平衡节点的右节点

					while (src&&src->_parent != unbalanced_node_right) {
						src = src->_parent;
					}
					if (unbalanced_node_right->_right == src) { //RR型 ：左单旋
						Rotation_Left(unbalanced_node, unbalanced_node_right);
					}
					else {  //RL型：
						Rotation_RL(unbalanced_node, unbalanced_node_right);
					}
				}
				else { //parent->_bf==-2
					TNode* unbalanced_node = parent; //不平衡节点
					TNode* unbalanced_node_left = parent->_left; //不平衡节点的左节点
					while (src&&src->_parent != unbalanced_node_left) {
						src = src->_parent;
					}
					if (unbalanced_node_left->_left == src) { //LL型：右单旋
						Rotation_Right(unbalanced_node, unbalanced_node_left);
					}
					else { //LR型：
						Rotation_LR(unbalanced_node, unbalanced_node_left);
					}
				}
				break;
			}
		  }
		}
	void InOrder() {
		Inorder(_root);
	}
	bool IsBalanceTree() {
		return _IsBalanceTree(_root);
	}
	int Height() {
		return TreeHeight(_root);
	}
	int Size() {
		return _Size(_root);
	}
	void Preorder() {
		_Preorder(_root);
	}
	~AVLTree() {
		DelAVLTree(_root);
		_root = nullptr;
	}
private:
	void DelAVLTree(TNode* root) {
		if (root == nullptr) {
			return;
		}
		DelAVLTree(root->_left);
		DelAVLTree(root->_right);
		delete root;
	}
	int _Size(TNode*root) {
		if (root == nullptr) {
			return 0;
		}
		return _Size(root->_left) + _Size(root->_right) + 1;
	}
	bool _IsBalanceTree(TNode*root) {
		if (root == nullptr) {
			return true;
		}
		int left_height = TreeHeight(root->_left);
		int right_height = TreeHeight(root->_right);

		int height = right_height - left_height;
		if (abs(height) >= 2) {
			cout << "该树不平衡" << endl;
			return false;
		}

		if (height != root->_bf) {
			cout <<root->_kv.first<< "  的平衡因子异常" << endl;
			return false;
		}

		//左右子树也必须平衡：
		return _IsBalanceTree(root->_left) && _IsBalanceTree(root->_right);
	}
	int TreeHeight(TNode* root) {
		if (root == nullptr) {
			return 0;
		}

		int left = TreeHeight(root->_left);
		int right = TreeHeight(root->_right);

		return left > right ? left + 1 : right + 1;
	}
	void _Preorder(TNode* root) {
		if (root == nullptr) {
			return;
		}
		cout << root->_kv.first<<" ";
		Inorder(root->_left);
		Inorder(root->_right);
	}
	void Inorder(TNode* root) {
		if (root == nullptr) {
			return;
		}
		Inorder(root->_left);
		cout << root->_kv.first <<" ";
		Inorder(root->_right);
	}
private:
	TNode* _root = nullptr;
};
